package com.woniuxy.handler;

import com.alibaba.fastjson.JSON;
import com.woniuxy.anno.RateLimit;
import com.woniuxy.util.ResponseData;
import com.woniuxy.util.ResponseEnum;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.stereotype.Component;
import org.springframework.web.method.HandlerMethod;
import org.springframework.web.servlet.HandlerInterceptor;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.lang.reflect.Method;
import java.util.concurrent.TimeUnit;

@Component
public class RateLimitInterceptor implements HandlerInterceptor {

    @Autowired
    private RedisTemplate<String, Long> redisTemplate;

    @Override
    public boolean preHandle(HttpServletRequest request, HttpServletResponse response, Object handler) throws Exception {
        if (handler instanceof HandlerMethod) {
            HandlerMethod handlerMethod = (HandlerMethod) handler;
            Method method = handlerMethod.getMethod();
            RateLimit rateLimitAnnotation = method.getAnnotation(RateLimit.class);
            if (rateLimitAnnotation != null) {
                int time = rateLimitAnnotation.time();
                int maxCount = rateLimitAnnotation.maxCount();
                String key = generateKey(request, method);
                // 在 Redis 中设置计数器，如果存在则递增，如果不存在则创建并设置过期时间
                Long count = increment(key, 1);
                if (count == 1) {
                    redisTemplate.expire(key, time, TimeUnit.SECONDS);
                }
                if (count != null && count > maxCount) {
                    // 达到最大连接次数，返回错误提示
                    response.setContentType("application/json;charset=UTF-8");
                    response.getWriter().write(JSON.toJSONString(new ResponseData().fail(ResponseEnum.RATE_LIMIT)));
                    return false;
                }
            }
        }
        return true;
    }

    private String generateKey(HttpServletRequest request, Method method) {
        // 获取请求的 IP 地址作为键
        String ipAddress = request.getRemoteAddr();
        String uri = request.getRequestURI();//请求的uri
        return "rate_limit:" + ipAddress + uri + ":" + method.getName();
    }

    public Long increment(String key, long delta) {
        // 递增操作，返回当前计数器的值
        return redisTemplate.opsForValue().increment(key, delta);
    }
}
